package com.breeze.components.web.filter;

import cn.hutool.core.io.IoUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HtmlUtil;
import cn.hutool.json.JSONUtil;
import jakarta.servlet.ReadListener;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.LinkedHashMap;
import java.util.Map;
import org.springframework.web.util.ContentCachingRequestWrapper;

/**
 * @author breeze
 */
public class XssCleanRequestWrapper extends HttpServletRequestWrapper {

  ByteArrayInputStream bain = null;

  /**
   * Constructs a request object wrapping the given request.
   *
   * @param request The request to wrap
   * @throws IllegalArgumentException if the request is null
   */
  public XssCleanRequestWrapper(HttpServletRequest request) {
    super(request);
  }

  // 提取请求参数
  @Override
  public String getParameter(String name) {
    String value = super.getParameter(name);
    if (!StrUtil.hasEmpty(value)) {
      value = HtmlUtil.escape(value);
    }
    return value;
  }

  @Override
  public String[] getParameterValues(String name) {
    String[] values = super.getParameterValues(name);
    if (values != null) {
      for (int i = 0; i < values.length; i++) {
        String value = values[i];
        if (!StrUtil.hasEmpty(value)) {
          value = HtmlUtil.escape(value);
        }
        values[i] = value;
      }
    }
    return values;
  }

  // 转义数据
  @Override
  public Map<String, String[]> getParameterMap() {
    Map<String, String[]> parameters = super.getParameterMap();
    LinkedHashMap<String, String[]> map = new LinkedHashMap<>();
    if (parameters != null) {
      for (String key : parameters.keySet()) {
        String[] values = parameters.get(key);
        for (int i = 0; i < values.length; i++) {
          String value = values[i];
          if (!StrUtil.hasEmpty(value)) {
            // 导入Hutool的依赖,进行html转义
            value = HtmlUtil.escape(value);
          }
          values[i] = value;
        }
        map.put(key, values);
      }
    }
    return map;
  }

  @Override
  public String getHeader(String name) {
    String value = super.getHeader(name);
    if (!StrUtil.hasEmpty(value)) {
      value = HtmlUtil.escape(value);
    }
    return value;
  }

  @Override
  public ServletInputStream getInputStream() throws IOException {
    String body;
    if (bain != null && this.getRequest() instanceof ContentCachingRequestWrapper request) {
      body = new String(request.getContentAsByteArray(), StandardCharsets.UTF_8);
    } else {
      bain = new ByteArrayInputStream(IoUtil.readBytes(super.getInputStream()));
      body = IoUtil.read(bain, StandardCharsets.UTF_8);
    }
    if (StrUtil.isNotEmpty(body)) {
      Map<String, Object> map = JSONUtil.parseObj(body);
      Map<String, Object> result = new LinkedHashMap<>();
      for (String key : map.keySet()) {
        Object val = map.get(key);
        if (val instanceof String) {
          if (!StrUtil.hasEmpty(val.toString())) {
            result.put(key, HtmlUtil.escape(val.toString()));
          }
        } else {
          result.put(key, val);
        }
      }
      bain = new ByteArrayInputStream(JSONUtil.toJsonStr(result).getBytes());
    } else {
      bain = new ByteArrayInputStream(StrUtil.EMPTY.getBytes());
    }
    // 匿名类实现IO流
    return new ServletInputStream() {
      @Override
      public int read() throws IOException {
        return bain.read();
      }

      @Override
      public boolean isFinished() {
        return false;
      }

      @Override
      public boolean isReady() {
        return false;
      }

      @Override
      public void setReadListener(ReadListener readListener) {}
    };
  }
}
